#pragma once

#include "../jsonxx/json.hpp"
#include "lexer_typedef_min.hpp"
#include <algorithm>
#include <deque>
#include <exception>
#include <iomanip>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <regex>

using namespace jsonxx;
using std::string;
using std::unordered_map;
using std::unordered_set;
using std::cerr;
using std::cout;
using std::deque;
using std::endl;
using std::fstream;
using std::make_shared;
using std::map;
using std::set;
using std::shared_ptr;
using std::stack;
using std::string;
using std::stringstream;
using std::unordered_map;
using std::unordered_set;
using std::vector;

template <typename Container_T, typename T>
int indexOf(Container_T &container, T &value) {
    for (int i = 0; i < container.size(); ++i) {
        if (container[i] == value)
            return i;
    }
    return -1;
}

string trim(const string &s) {
    if (s.empty())
        return s;
    int stPos = s.find_first_not_of(" ");
    if (stPos == string::npos)
        return "";
    return s.substr(stPos, s.find_last_not_of(" ") + 1);
}

string replace(const string &str, const string &reg, const string &target) {
    return std::regex_replace(str, std::regex(reg), target);
}

/**
 * if tokenCompress is set to true, then continuous delimiter would treated as one
 */
std::vector<std::string> split(const std::string &in, const std::string &regStr,
                               bool tokenCompress = true) {
    std::regex reg(regStr);
    std::vector<std::string> v(std::sregex_token_iterator(in.begin(), in.end(), reg, -1),
                               std::sregex_token_iterator());
    if (!tokenCompress)
        return v;
    std::vector<std::string> newV;
    for (auto i : v) {
        if (!i.empty()) {
            newV.push_back(i);
        }
    }
    return newV;
}

class MyParser {
public:
    enum class ActionType { //分析表动作类型
        UNDEFINED,
        ACCEPT,
        SHIFT,
        REDUCE,
        CONFLICT,
        ERROR
    };
    class RuleType;      // 规则类型
    class Action;        // 分析表动作
    class ParseTreeNode; // 分析树节点
    using PtrAction = shared_ptr<Action>;
    using PtrParseTreeNode = shared_ptr<ParseTreeNode>;

public:
    MyParser(const vector<string> &symbols, const vector<vector<Action>> &parseTab)
        : _symbols(symbols), _parseTable(parseTab) {}
    void createParser(std::istream &iss);
    PtrParseTreeNode parse(const vector<Token> &tokens);
    void printParseTable(std::ostream &oss);
    void printParseTree(std::ostream &oss, const PtrParseTreeNode &root, string indent);

public:
    /**
     * 规则类型
     */
    class RuleType {
    public:
        string name; // lhs在symbol里的名字
        int index;   // 该规则在rules数组中的下标
        int lhs;     // 该规则名称在symbols数组中的下标
        int rhsSize; // 该规则包含的子规则和终结符在symbols数组中的下标
        RuleType();
        RuleType(string name, int index, int lhs, int rhsSize)
            : name(name), index(index), lhs(lhs), rhsSize(rhsSize){}
        string getName() const;
    };

    /**
     * 分析表动作
     */
    class Action {
    public:
        ActionType atype;
        int state;
        RuleType rule;
        Action();
        Action(bool);
        Action(int state);
        Action(RuleType rule);
        ~Action() {}
        friend std::ostream &operator<<(std::ostream &oss, const Action &a);
    };

    /**
     * 分析树节点
     */
    class ParseTreeNode {
    public:
        union {
            RuleType ruleT;
            Token token;
        };
        deque<PtrParseTreeNode> children;

        bool isTerminal();
        ParseTreeNode();
        ParseTreeNode(RuleType ruleT, deque<PtrParseTreeNode> children);
        ParseTreeNode(Token token);
        ~ParseTreeNode();
    };

private:
    const vector<string> _symbols;    // 语法符号
    const vector<vector<Action>> _parseTable;          // 动作
    unordered_map<string, int> _tokenName2Index; // 将token名转换成index
    int token2Symbol(const Token &token);
};

int MyParser::token2Symbol(const Token &token) {
    string tokenName = token.tokenName;
    if (_tokenName2Index.count(tokenName)) {
        return _tokenName2Index.at(tokenName);
    } else {
        int index = indexOf(_symbols, tokenName);
        if (index == -1) {
            cout << "[Error] Token \'" << tokenName
                 << "\' did not exist in grammar definition, "
                 << "please check your grammar definition file.\n"
                 << "Parser terminated.\n";
            throw ParseException(token.endPos.line, token.endPos.col, token.word);
            exit(0);
        }
        _tokenName2Index[tokenName] = index;
        return index;
    }
}

MyParser::PtrParseTreeNode MyParser::parse(const vector<Token> &input) {
    stack<int> stateStack;
    stack<PtrParseTreeNode> symbolStack;
    Action action;
    stateStack.push(0);
    vector<Token>::const_iterator inputIter = input.begin();
    while (true) {
        if (inputIter == input.end()) {
            action = _parseTable[stateStack.top()][_symbols.size() - 1];
        } else {
            action = _parseTable[stateStack.top()][token2Symbol(*inputIter)];
        }
        if (action.atype == ActionType::UNDEFINED) {
            cout << "[Error] line " << inputIter->startPos.line << "("
                 << inputIter->startPos.col << "): \'" << inputIter->word
                 << "\'\nSyntax error!";
            throw ParseException(inputIter->startPos.line, inputIter->startPos.col,
                                 inputIter->word);
            exit(0);
        }
        if (action.atype == ActionType::SHIFT) {
            symbolStack.push(make_shared<ParseTreeNode>(*inputIter));
            inputIter++;
            stateStack.push(action.state);
        } else if (action.atype == ActionType::REDUCE) {
            deque<PtrParseTreeNode> children; // 如果规约，则把所有归约的字符加进这里
            for (int i = 0; i < action.rule.rhsSize; ++i) {
                children.push_front(symbolStack.top());
                symbolStack.pop();
                stateStack.pop();
            }
            PtrParseTreeNode node = make_shared<ParseTreeNode>(action.rule, children);
            symbolStack.push(node);
            // GOTO
            stateStack.push(_parseTable[stateStack.top()][action.rule.lhs].state);
        } else if (action.atype == ActionType::ACCEPT) {
            break;
        } else {
            cout << "Syntax error" << endl;
            exit(0);
        }
    }
    return symbolStack.top();
}

void MyParser::printParseTree(std::ostream &oss, const MyParser::PtrParseTreeNode &t,
                              string indent = "") {
    oss << indent << "|-- ";
    if (t->isTerminal()) {
        oss << t->token.word << endl;
    } else {
        oss << _symbols[t->ruleT.lhs] << endl;
        for (auto i : t->children) {
            printParseTree(oss, i, indent + "  ");
        }
    }
}

void MyParser::printParseTable(std::ostream &grammarOut) {
    grammarOut << "state, ";
    for (int i = 1; i < _symbols.size(); ++i) {
        grammarOut << _symbols[i] << ", ";
    }
    grammarOut << endl;
    for (int i = 0; i < _parseTable.size(); ++i) {
        grammarOut << i << ",";
        for (int j = 1; j < _parseTable[i].size(); ++j) {
            grammarOut << _parseTable[i][j] << ", ";
        }
        grammarOut << endl;
    }
}

MyParser::RuleType::RuleType(): name(""){};
string MyParser::RuleType::getName() const { return name; }


MyParser::Action::Action() : atype(ActionType::UNDEFINED) {}
MyParser::Action::Action(bool) : atype(ActionType::ACCEPT) {}
MyParser::Action::Action(int state) : atype(ActionType::SHIFT), state(state) {}
MyParser::Action::Action(RuleType rule) : atype(ActionType::REDUCE), rule(rule) {}
std::ostream &operator<<(std::ostream &oss, const MyParser::Action &a) {
    switch (a.atype) {
    case MyParser::ActionType::UNDEFINED:
        oss << "-";
        break;
    case MyParser::ActionType::SHIFT:
        oss << "s" << a.state;
        break;
    case MyParser::ActionType::REDUCE:
        oss << "r" << a.rule.index;
        break;
    case MyParser::ActionType::ACCEPT:
        oss << "ACC";
        break;
    default:
        oss << __FILE__ << ": " << __LINE__ << "wtf";
        exit(0);
    }
    return oss;
}


bool MyParser::ParseTreeNode::isTerminal() { return children.empty(); }
MyParser::ParseTreeNode::ParseTreeNode() {}
MyParser::ParseTreeNode::ParseTreeNode(RuleType ruleT, deque<PtrParseTreeNode> children)
    : ruleT(ruleT), children(children){};
MyParser::ParseTreeNode::ParseTreeNode(Token token) : token(token){};
MyParser::ParseTreeNode::~ParseTreeNode() {}

class ASTConfig {
public:
    unordered_set<string> ignoredTerminals;
    unordered_set<string> childToSibling;
    unordered_map<string, string> generalization;
    unordered_set<string> compressExclusion;
    unordered_map<string, string> terminalNameMap;
    unordered_set<string> childToFields;
    unordered_set<string> mergeChilds;
    /**
     * Factory method to build ASTConfig from input stream
     * @param fin input stream
     * @return ASTConfig created from the json config
     */
    static ASTConfig configWithJson(std::istream &fin) {
        ASTConfig config;
        json jd;
        fin >> jd;
        auto &&ignore = jd["ignoredTerminals"].as_array();
        config.ignoredTerminals = {ignore.begin(), ignore.end()};

        auto &&c2f = jd["childToFields"].as_array();
        config.childToFields = {c2f.begin(), c2f.end()};

        auto &&c2s = jd["childToSiblings"].as_array();
        config.childToSibling = {c2s.begin(), c2s.end()};

        auto &&mc = jd["mergeChilds"].as_array();
        config.mergeChilds = {mc.begin(), mc.end()};

        auto &&nameMap = jd["terminalNamesMapping"];
        for (auto iter = nameMap.begin(); iter != nameMap.end(); iter++) {
            for (auto terminal = iter.value().begin(); terminal != iter.value().end();
                 ++terminal) {
                config.terminalNameMap.insert({terminal.value().as_string(), iter.key()});
            }
        }

        auto &&rules = jd["parseTreeToASTRules"].as_array();
        for (auto i : rules) {
            if (i["type"] == "Generalization") {
                auto targets = i["targets"];
                for (auto j : targets) {
                    config.generalization.insert(
                        {j.as_string(), i["category"].as_string()});
                }
                if (i.count("compressExclusions")) {
                    config.compressExclusion = {
                        i["compressExclusions"].as_array().begin(),
                        i["compressExclusions"].as_array().end()};
                }
            }
        }
        std::cout << "Initialized AST static fields\n";
        return config;
    }
};

class ASTNode;

using PtrASTNode = shared_ptr<ASTNode>;

class ASTNode {
public:
    string name;
    unordered_map<string, vector<Token>> fields;
    vector<PtrASTNode> children;
    ASTNode() = default;
    void printAST(std::ostream &oss) { ASTNode::printAST(oss, this); }
    static PtrASTNode buildAST(MyParser::PtrParseTreeNode p, const ASTConfig &config) {
        PtrASTNode ast = buildASTStep1_rawAST(p, config);
        cout << "[Info] Finish step 1\n";
        ASTNode::buildASTStep2_child2Field(ast, config);
        cout << "[Info] Finish step 2\n";
        ASTNode::buildASTStep3_child2Sibling(ast, config);
        cout << "[Info] Finish step 3\n";
        ASTNode::buildASTStep4_mergeChilds(ast, config);
        cout << "[Info] Finish step 4\n";
        cout << "[Info] Finish build AST;\n";
        return ast;
    }

private:
    template <class Ptr>
    static void printAST(std::ostream &oss, const Ptr &t, const string &indent = "") {
        oss << indent << "|-- " << t->name << "{ ";
        int cnt = 0;
        for (auto i : t->fields) {
            oss << i.first << ": ";
            if (i.second.empty())
                oss << i.second[0].word;
            else {
                for (auto j : i.second) {
                    oss << j.word << " ";
                }
            }
            ++cnt;
            if (cnt < t->fields.size())
                oss << ", ";
        }
        oss << " }\n";
        for (auto i : t->children) {
            printAST(oss, i, indent + "  ");
        }
    }

private:
    static PtrASTNode buildASTStep1_rawAST(MyParser::PtrParseTreeNode p,
                                           const ASTConfig &config) { // 第一步，建立AST
        PtrASTNode ast = make_shared<ASTNode>();
        vector<PtrASTNode> newChildren;
        ast->name = p->ruleT.getName();
        for (auto i : p->children) {
            if (!i->isTerminal()) { // 递归
                newChildren.push_back(buildASTStep1_rawAST(i, config));
            } else if (!config.ignoredTerminals.count(i->token.tokenName)) {
                // 把不可删除的终结符提升到树根
                if (ast->fields.count(config.terminalNameMap.at(i->token.tokenName))) {
                    ast->fields.at(config.terminalNameMap.at(i->token.tokenName))
                        .push_back(i->token);
                } else {
                    ast->fields[config.terminalNameMap.at(i->token.tokenName)] = {
                        i->token};
                }
            }
        }
        if (config.generalization.count(p->ruleT.getName())) {
            if (newChildren.size() == 1 &&
                !config.compressExclusion.count(p->ruleT.getName())) {
                return newChildren[0];
            }
            ast->name = config.generalization.at(p->ruleT.getName());
        }
        ast->children = newChildren;
        return ast;
    }

    static void
    buildASTStep2_child2Field(PtrASTNode ast,
                              const ASTConfig &config) { // 第二步，child to fields
        vector<PtrASTNode> newChildren;
        for (auto i : ast->children) {
            if (i->children.empty() && config.childToFields.count(i->name)) { // 压缩
                ast->fields.insert(i->fields.begin(), i->fields.end());
            } else {
                newChildren.push_back(i);
                buildASTStep2_child2Field(i, config);
            }
        }
        ast->children = newChildren;
    }

    static vector<PtrASTNode>
    buildASTStep3_child2Sibling(PtrASTNode ast,
                                const ASTConfig &config) { // 第三步，child to sibling
        vector<PtrASTNode> other2bElevate; // 其它类型递归节点的养父
        vector<PtrASTNode> self2bElevate;  // 自己类型递归节点的养父
        vector<PtrASTNode> newChildren;
        unordered_set<string> recursiveChildNames;
        bool hasRecursiveChild = false;
        for (auto i : ast->children) {
            auto &&tmp = buildASTStep3_child2Sibling(i, config); // 获得孩子的养父
            if (config.childToSibling.count(i->name)) {          // 是要被提升的
                hasRecursiveChild = true;
                recursiveChildNames.insert(i->name);
                if (ast->name != i->name) {
                    other2bElevate.insert(other2bElevate.end(), tmp.begin(), tmp.end());
                } else {
                    self2bElevate.insert(self2bElevate.end(), tmp.begin(), tmp.end());
                }
            } else {
                newChildren.push_back(i);
            }
        }
        newChildren.insert(newChildren.end(), other2bElevate.rbegin(),
                           other2bElevate.rend());
        ast->children = newChildren;
        if (config.childToSibling.count(ast->name)) { // 自己就是在递归
            self2bElevate.push_back(ast);
            return self2bElevate;
        }
        return {};
    }

    static void buildASTStep4_mergeChilds(PtrASTNode ast, const ASTConfig &config) {
        unordered_map<string, PtrASTNode> stepFathers; // <节点类型，节点的孩子>
        vector<PtrASTNode> newChildren;
        for (auto i : ast->children) {
            buildASTStep4_mergeChilds(i, config);
            if (config.mergeChilds.count(i->name)) {
                if (!stepFathers.count(i->name)) {
                    stepFathers[i->name] = make_shared<ASTNode>();
                }
                stepFathers[i->name]->children.insert(
                    stepFathers[i->name]->children.end(), i->children.begin(),
                    i->children.end());
            } else {
                newChildren.push_back(i);
            }
        }
        for (auto [name, father] : stepFathers) {
            father->name = name;
            newChildren.push_back(father);
        }
        ast->children = newChildren;
    }
};

template <> struct json_bind<PtrASTNode> {
    const string NAME = "Name";
    const string CHILDREN = "children";
    void to_json(json &j, const PtrASTNode &ast) {
        jsonxx::to_json(j[NAME], ast->name);
        for (auto [key, val] : ast->fields) {
            string str;
            for (auto j : val)
                str += j.word;
            jsonxx::to_json(j[key], str);
        }
        vector<json> children;
        for (auto i = ast->children.begin(); i != ast->children.end(); ++i) {
            json tmp;
            to_json(tmp, *i);
            children.push_back(tmp);
        }
        if (!children.empty())
            j[CHILDREN] = children;
    }

    void from_json(const json &j, PtrASTNode &ast) {
        jsonxx::from_json(j[NAME], ast->name);
        vector<json> children;
        if (j.count(CHILDREN))
            children = j[CHILDREN].as_array();
        for (auto i : children) {
            PtrASTNode tmp = make_shared<ASTNode>();
            from_json(i, tmp);
            ast->children.push_back(tmp);
        }
        for (auto [key, val] : j.as_object()) {
            if (key != CHILDREN) { // 处理fields
                Token token;
                token.word = val.as_string();
                ast->fields[key].push_back(token);
            }
        }
    }
};
